[model] feat: add MiniMax-M2 MoE bridge - #2602

tests/unit_tests/recipes/test_moonlight_recipes.py (1)

18-18: Consider updating stale documentation and class naming.

The docstring on line 18 still references MoonlightModelProvider16B, and the fake class _FakeMoonlightModelProvider16B (line 76) is now used to mock MLAModelProvider. This inconsistency could confuse future maintainers.

✏️ Suggested documentation updates

 # Test purpose:
 # - Parametrize over all exported Moonlight recipe functions in `megatron.bridge.recipes.moonlight`.
-# - For each recipe, monkeypatch `MoonlightModelProvider16B` with a lightweight fake to avoid I/O.
+# - For each recipe, monkeypatch `MLAModelProvider` with a lightweight fake to avoid I/O.
 # - Build a config with small, safe overrides and assert it forms a valid `ConfigContainer`.
 # - Verify tokenizer selection and sanity-check parallelism fields.

-class _FakeMoonlightModelProvider16B:
-    """Fake MoonlightModelProvider16B for testing without model I/O."""
+class _FakeMLAModelProvider:
+    """Fake MLAModelProvider for testing Moonlight recipes without model I/O."""

If renaming the class, update all usages (lines 136, 171, 196, 222, 238, 261, 284, 307, 329).

Also applies to: 76-77

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unit_tests/recipes/test_moonlight_recipes.py` at line 18, Update the
stale docstring and fake class name so they match the current provider type:
change any docstring text referencing MoonlightModelProvider16B to reference
MLAModelProvider (or the current provider name) and rename the fake class
_FakeMoonlightModelProvider16B to _FakeMLAModelProvider (or similar) and update
all monkeypatch usages that reference the old name (the test class and each
place where the fake is injected in the tests). Ensure the docstring, the fake
class declaration, and every usage/monkeypatch that currently points to
_FakeMoonlightModelProvider16B are updated to the new MLAModelProvider-based
names so they are consistent.

tests/unit_tests/training/test_config.py (1)

79-89: Consider renaming this function to reflect its current purpose.

The function is named create_test_deepseek_config but now creates an MLAModelProvider. Consider renaming to create_test_mla_config for consistency with the updated implementation. The test comment on line 1229 also references "Deepseek" but tests MLAModelProvider.

♻️ Suggested rename

-def create_test_deepseek_config(**kwargs: Any) -> MLAModelProvider:
-    """Creates an instance of MLAModelProvider for testing."""
+def create_test_mla_config(**kwargs: Any) -> MLAModelProvider:
+    """Creates an instance of MLAModelProvider for testing MLA-based models."""
     defaults = {
         "num_layers": 1,
         "hidden_size": 128,
         "num_attention_heads": 4,
         "seq_length": 512,
         "apply_rope_fusion": False,
     }
     defaults.update(kwargs)
     return MLAModelProvider(**defaults)

Also update the reference on line 1225 and comment on line 1229:

-    `@pytest.mark.parametrize`("model_factory", [create_test_gpt_config, create_test_deepseek_config])
+    `@pytest.mark.parametrize`("model_factory", [create_test_gpt_config, create_test_mla_config])
     def test_default_pipeline_dtype(self, model_factory, monkeypatch):
         """
-        Test pipeline_dtype is automatically set if None and PP enabled.
-        Test for both GPT and Deepseek to test both TransformerConfig types.
+        Test pipeline_dtype is automatically set if None and PP enabled.
+        Test for both GPT and MLA to test both TransformerConfig types.
         """

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unit_tests/training/test_config.py` around lines 79 - 89, Rename the
helper function create_test_deepseek_config to create_test_mla_config and update
all call sites and comments that still reference "Deepseek" to instead reference
MLAModelProvider/MLA (e.g., change the function name and any test references and
the comment that mentions Deepseek to reflect MLAModelProvider), ensuring the
factory still returns MLAModelProvider(**defaults) and preserving the same
default kwargs.

tests/unit_tests/training/test_log_non_default_values.py (2)

58-71: Test logic is correct; consider renaming for clarity.

The test correctly verifies that MLAModelProvider returns MCoreMLATransformerConfig as its parent, and the docstring is accurate. The method name test_deepseek_provider_returns_mla_transformer_config could be renamed to test_mla_provider_returns_mla_transformer_config for consistency.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unit_tests/training/test_log_non_default_values.py` around lines 58 -
71, Rename the test function to match the subject under test: change the
function name test_deepseek_provider_returns_mla_transformer_config to
test_mla_provider_returns_mla_transformer_config in the test file so it clearly
reflects MLAModelProvider; keep the existing docstring and assertions (including
the reference to _get_mcore_transformer_parent and MCoreMLATransformerConfig)
unchanged.

---

299-315: Test logic is correct; naming could be improved for consistency.

The test correctly verifies that MLAModelProvider uses MLATransformerConfig for comparison. However, the method name, docstring, and variable name still reference "DeepSeek" while testing the generic MLAModelProvider.

♻️ Suggested naming updates

     `@patch`("megatron.bridge.training.config.print_rank_0")
-    def test_handles_deepseek_model_correctly(self, mock_print_rank_0):
-        """Should use MLATransformerConfig for DeepSeek models."""
-        deepseek_model = MLAModelProvider(
+    def test_handles_mla_model_correctly(self, mock_print_rank_0):
+        """Should use MLATransformerConfig for MLA-based models."""
+        mla_model = MLAModelProvider(
             num_layers=2,
             hidden_size=128,
             num_attention_heads=4,
             seq_length=512,
         )

-        cfg = self._create_minimal_config_container(model_provider=deepseek_model)
+        cfg = self._create_minimal_config_container(model_provider=mla_model)

         cfg.log_non_default_values()

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unit_tests/training/test_log_non_default_values.py` around lines 299 -
315, Rename references to "DeepSeek" in the test to reflect the generic
MLAModelProvider: update the test method name
test_handles_deepseek_model_correctly, its docstring, and the variable
deepseek_model to names that reference MLAModelProvider (e.g.,
test_handles_mla_model_correctly and mla_model) so they match the actual class
under test (MLAModelProvider); keep the assertions and usage of
mock_print_rank_0, cfg, and cfg.log_non_default_values() unchanged and ensure
the assertion still verifies "MLATransformerConfig" appears in
mock_print_rank_0.call_args[0][0].

examples/conversion/compare_hf_and_megatron/compare.py (1)

612-613: Also disable grad_scale_func when turning off MTP.

For consistency with the existing MTP-disable pattern, set m.config.grad_scale_func = None alongside mtp_num_layers = None.

🔧 Suggested patch

     for m in megatron_model:
         m.config.mtp_num_layers = None
+        m.config.grad_scale_func = None

Based on learnings: In examples/conversion/compare_hf_and_megatron/compare.py, disabling MTP for inference used m.config.mtp_num_layers = None and m.config.grad_scale_func = None together.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@examples/conversion/compare_hf_and_megatron/compare.py` around lines 612 -
613, When disabling MTP in the loop over megatron_model, also set the grad scale
function to None: inside the for m in megatron_model loop where you currently do
m.config.mtp_num_layers = None, add m.config.grad_scale_func = None so both
m.config.mtp_num_layers and m.config.grad_scale_func are cleared to disable MTP
consistently for inference.

src/megatron/bridge/recipes/moonlight/moonlight_16b.py (1)

199-255: Deduplicate shared MLAModelProvider kwargs between SFT and PEFT.

These two constructor blocks are large and mostly identical, which increases config drift risk.

♻️ Suggested refactor sketch

+def _moonlight_16b_model_kwargs(*, tp: int, ep: int, sequence_parallel: bool) -> dict:
+    return dict(
+        num_layers=27,
+        hidden_size=2048,
+        ffn_hidden_size=11264,
+        num_attention_heads=16,
+        kv_channels=16,
+        q_lora_rank=None,
+        kv_lora_rank=512,
+        num_moe_experts=64,
+        moe_ffn_hidden_size=1408,
+        moe_shared_expert_intermediate_size=2816,
+        moe_layer_freq=[0] * 1 + [1] * 26,
+        moe_router_topk=6,
+        moe_router_num_groups=1,
+        moe_router_group_topk=1,
+        moe_router_topk_scaling_factor=2.446,
+        moe_aux_loss_coeff=0.001,
+        make_vocab_size_divisible_by=1280,
+        moe_router_score_function="sigmoid",
+        moe_router_enable_expert_bias=True,
+        rotary_scaling_factor=1.0,
+        mscale=1.0,
+        mscale_all_dim=1.0,
+        rotary_base=50000,
+        layernorm_epsilon=1e-5,
+        init_method_std=0.02,
+        moe_router_bias_update_rate=1e-3,
+        rotary_percent=1.0,
+        vocab_size=163842,
+        normalization="RMSNorm",
+        activation_func=F.silu,
+        gated_linear_unit=True,
+        position_embedding_type="rope",
+        add_bias_linear=False,
+        share_embeddings_and_output_weights=False,
+        qk_layernorm=True,
+        bf16=True,
+        params_dtype=torch.bfloat16,
+        moe_grouped_gemm=True,
+        moe_token_dispatcher_type="alltoall",
+        tensor_model_parallel_size=tp,
+        pipeline_model_parallel_size=1,
+        pipeline_dtype=torch.bfloat16,
+        virtual_pipeline_model_parallel_size=None,
+        context_parallel_size=1,
+        expert_model_parallel_size=ep,
+        sequence_parallel=sequence_parallel,
+        expert_tensor_parallel_size=1,
+        recompute_granularity="selective",
+        recompute_modules=None,
+        recompute_method=None,
+        recompute_num_layers=None,
+    )
@@
-    cfg.model = MLAModelProvider(
-        ...
-    )
+    cfg.model = MLAModelProvider(**_moonlight_16b_model_kwargs(tp=2, ep=8, sequence_parallel=True))
@@
-    cfg.model = MLAModelProvider(
-        ...
-    )
+    cfg.model = MLAModelProvider(**_moonlight_16b_model_kwargs(tp=1, ep=2, sequence_parallel=False))

Also applies to: 413-468

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/recipes/moonlight/moonlight_16b.py` around lines 199 -
255, The MLAModelProvider constructor kwargs are duplicated between the SFT and
PEFT blocks (MLAModelProvider in moonlight_16b.py); extract the shared keyword
args into a single dict (e.g., common_model_kwargs or common_mla_config) and
pass it into both MLAModelProvider(...) calls (merge per-block overrides after
the shared dict), so SFT and PEFT only specify differences. Update any
references to parameters like num_layers, hidden_size, moe_*, rotary_*,
normalization, activation_func, bf16, params_dtype, parallelism and recompute_*
to live in the shared dict and leave only unique overrides in the individual
SFT/PEFT constructor calls.

tests/unit_tests/recipes/test_glm45_recipes.py (1)

138-153: Extract common monkeypatch setup into a pytest fixture.

This setup pattern is repeated and should be centralized to reduce drift across tests.

♻️ Example fixture extraction

+@pytest.fixture
+def patch_glm_recipe_io(monkeypatch: pytest.MonkeyPatch):
+    def _apply(mod, needs_tokenizer: bool):
+        monkeypatch.setattr(mod, "AutoBridge", _FakeBridge)
+        if needs_tokenizer:
+            import transformers
+            monkeypatch.setattr(
+                transformers,
+                "AutoTokenizer",
+                type("FakeAutoTokenizer", (), {"from_pretrained": staticmethod(lambda *args, **kwargs: _FakeTokenizer())}),
+            )
+    return _apply
@@
-def test_each_glm45_recipe_builds_config(recipe_func: Callable, monkeypatch: pytest.MonkeyPatch):
+def test_each_glm45_recipe_builds_config(recipe_func: Callable, patch_glm_recipe_io):
@@
-    monkeypatch.setattr(mod, "AutoBridge", _FakeBridge)
+    patch_glm_recipe_io(mod, is_sft_or_peft)

As per coding guidelines, "tests/**/*.py: Use pytest fixtures for common setup in unit tests."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/unit_tests/recipes/test_glm45_recipes.py` around lines 138 - 153,
Extract the repeated monkeypatch setup into a pytest fixture (e.g.,
fake_bridge_and_tokenizer) that accepts monkeypatch and recipe_func (or uses
request.param) and performs the same work: import the recipe module via
recipe_func.__module__, monkeypatch the module's AutoBridge to _FakeBridge,
compute is_sft_or_peft by checking "sft" or "peft" in
recipe_func.__name__.lower(), and when true monkeypatch
transformers.AutoTokenizer to a FakeAutoTokenizer that returns
_FakeTokenizer.from_pretrained; replace the inline code in tests with a use of
this fixture to centralize the AutoBridge/AutoTokenizer mocking and avoid
duplication.

examples/models/minimax_m2/inference.sh (1)

1-24: Add strict shell mode for safer multi-node execution.

Please add fail-fast shell options so partial failures do not continue silently.

♻️ Suggested hardening

 #!/usr/bin/env bash
+set -euo pipefail
+
 # Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
@@
-uv run python -m torch.distributed.run --nproc_per_node=8 \
+uv run --no-sync python -m torch.distributed.run --nproc_per_node=8 \

As per coding guidelines, "**/*.sh: Follow Google Shell Style Guide." Based on learnings, "prefer including --no-sync` when dependencies are already installed and environment/toolchain compatibility is known."

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@examples/models/minimax_m2/inference.sh` around lines 1 - 24, Add fail-fast
shell options and the recommended --no-sync flag: at the top of inference.sh
(right after the shebang) enable strict mode by adding "set -euo pipefail"
(optionally "set -x" for debug) so the script exits on any error/unset
variable/pipe failure, and update the distributed run command line (the "uv run
python -m torch.distributed.run --nproc_per_node=8 ...") to include the
"--no-sync" option when appropriate to avoid unnecessary dependency sync; ensure
any future shell variables in this file are quoted to remain safe under "set
-u".

src/megatron/bridge/recipes/nemotronh/nemotron_nano_v2.py (1)

37-77: Extract a shared Nano-v2 provider factory to reduce drift risk.

These six MambaModelProvider blocks duplicate a large set of identical fields. A small helper with variant-specific overrides would make updates safer and easier.

♻️ Refactor direction (example)

+def _build_nemotron_nano_v2_provider(
+    *,
+    variant: str,  # "9b" | "12b"
+    seq_length: int,
+    tensor_model_parallel_size: int,
+    sequence_parallel: bool,
+) -> MambaModelProvider:
+    base_kwargs = dict(
+        mamba_num_groups=8,
+        num_query_groups=8,
+        make_vocab_size_divisible_by=128,
+        activation_func=squared_relu,
+        masked_softmax_fusion=True,
+        apply_query_key_layer_scaling=False,
+        persist_layer_norm=True,
+        attention_softmax_in_fp32=False,
+        first_last_layers_bf16=True,
+        is_hybrid_model=True,
+        moe_aux_loss_coeff=0.0001,
+        moe_router_score_function="sigmoid",
+        moe_router_enable_expert_bias=True,
+        moe_router_load_balancing_type="seq_aux_loss",
+        moe_router_dtype="fp32",
+        moe_grouped_gemm=True,
+        moe_token_dispatcher_type="alltoall",
+        moe_permute_fusion=True,
+        moe_shared_expert_overlap=True,
+        pipeline_model_parallel_size=1,
+        pipeline_dtype=torch.bfloat16,
+        virtual_pipeline_model_parallel_size=None,
+        context_parallel_size=1,
+    )
+    variant_kwargs = {
+        "9b": dict(...),
+        "12b": dict(...),
+    }[variant]
+    return MambaModelProvider(
+        **base_kwargs,
+        **variant_kwargs,
+        seq_length=seq_length,
+        tensor_model_parallel_size=tensor_model_parallel_size,
+        sequence_parallel=sequence_parallel,
+    )

Also applies to: 172-212, 309-349, 435-475, 571-610, 718-757

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@src/megatron/bridge/recipes/nemotronh/nemotron_nano_v2.py` around lines 37 -
77, Multiple identical MambaModelProvider blocks (e.g., the cfg.model
assignment) should be consolidated into a shared factory to avoid drift; create
a helper function (e.g., make_nano_v2_provider or create_nano_v2_provider) that
returns a MambaModelProvider pre-populated with the common fields shown
(hybrid_override_pattern, num_layers, hidden_size, mamba_num_heads, kv_channels,
mamba_state_dim, ffn_hidden_size, num_attention_heads, mamba_head_dim,
seq_length, mamba_num_groups, num_query_groups, make_vocab_size_divisible_by,
activation_func, masked_softmax_fusion, apply_query_key_layer_scaling,
persist_layer_norm, attention_softmax_in_fp32, first_last_layers_bf16,
is_hybrid_model, moe_* settings, parallelism defaults, etc.), then replace each
duplicated MambaModelProvider instantiation (including cfg.model and the other
five blocks) with a call to this factory passing only variant-specific overrides
(e.g., tensor_model_parallel_size, pipeline_model_parallel_size, pipeline_dtype,
virtual_pipeline_model_parallel_size, context_parallel_size, sequence_parallel)
so updates to common settings are made in one place.

tests/functional_tests/training/test_decentralized_pg.py (1)

104-140: Extract a shared model-config factory to reduce maintenance drift across duplicated test configurations.

The same GPTModelProvider baseline is repeated across six test functions. A local helper with base kwargs and per-test overrides would eliminate drift-prone duplication and make test-specific variations clearer.

♻️ Refactor sketch

+    def _create_model_cfg(self, **overrides):
+        base = dict(
+            normalization="RMSNorm",
+            activation_func=F.silu,
+            gated_linear_unit=True,
+            position_embedding_type="rope",
+            add_bias_linear=False,
+            attention_dropout=0.0,
+            hidden_dropout=0.0,
+            bias_activation_fusion=True,
+            masked_softmax_fusion=True,
+            persist_layer_norm=True,
+            bias_dropout_fusion=True,
+            apply_rope_fusion=True,
+            num_query_groups=8,
+            init_method_std=0.02,
+            layernorm_epsilon=1e-05,
+            rotary_percent=1.0,
+            rope_scaling=True,
+            rope_scaling_factor=32.0,
+            rotary_base=500_000,
+            hidden_size=2048,
+            ffn_hidden_size=8192,
+            num_attention_heads=32,
+            attention_softmax_in_fp32=True,
+            pipeline_dtype=torch.bfloat16,
+            bf16=True,
+            make_vocab_size_divisible_by=128,
+            vocab_size=None,
+        )
+        base.update(overrides)
+        return GPTModelProvider(**base)

Applies to lines 104–140, 255–290, 410–446, 566–602, 722–758, 878–914.

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@tests/functional_tests/training/test_decentralized_pg.py` around lines 104 -
140, The repeated GPTModelProvider instantiation should be refactored into a
shared factory: create a local helper (e.g., make_gpt_model_cfg or
gpt_model_provider_factory) that constructs the base kwargs shown
(normalization, activation_func, gated_linear_unit, position_embedding_type,
etc.), accepts overrides for per-test changes, and returns a GPTModelProvider
instance; replace the six duplicated blocks that assign model_cfg with calls to
this helper, passing only test-specific overrides (like seq_length, num_layers,
share_embeddings_and_output_weights), so maintenance drift is eliminated and
variations remain explicit.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro